package com.dataStructures.huffmancode;

import java.io.*;
import java.util.*;

public class HuffmanCode {
	public static void main(String[] args) {
		String zipFile = "/Users/zhangchuanqiang/Desktop/aa.zip";
		String dstFile = "/Users/zhangchuanqiang/Desktop/aa1.png";
		unZipFile(zipFile, dstFile);
//		zipFile(zipFile, dstFile);
		System.out.println("success");
	}

	/**
	 * 赫夫曼表
	 */
	static Map<Byte, String> huffmanCodes = new HashMap<>();
	static StringBuilder stringBuilder = new StringBuilder();

	/**
	 * 文件压缩
	 *
	 * @param srcFile 原文件
	 * @param dstFile 压缩文件
	 */
	public static void zipFile(String srcFile, String dstFile) {
		OutputStream os = null;
		ObjectOutputStream oos = null;
		FileInputStream is = null;
		try {
			is = new FileInputStream(srcFile);
			byte[] b = new byte[is.available()];
			is.read(b);
			byte[] huffmanBytes = huffmanZip(b);
			os = new FileOutputStream(dstFile);
			oos = new ObjectOutputStream(os);
			// 把压缩后的字节数组写入压缩文件
			oos.writeObject(huffmanBytes);
			// 把赫夫曼表写入压缩文件
			oos.writeObject(huffmanCodes);
		} catch (Exception e) {
			System.out.println(e.getMessage());
		} finally {
			try {
				is.close();
				oos.close();
				os.close();
			} catch (Exception e) {
				System.out.println(e.getMessage());
			}
		}
	}

	/**
	 * 文件解压缩
	 *
	 * @param zipFile 压缩文件
	 * @param dstFile 目标文件
	 */
	public static void unZipFile(String zipFile, String dstFile) {
		InputStream is = null;
		ObjectInputStream ois = null;
		OutputStream os = null;
		try {
			is = new FileInputStream(zipFile);
			ois = new ObjectInputStream(is);
			// 读取赫夫曼码表
			byte[] huffmanBytes = (byte[]) ois.readObject();
			Map<Byte, String> huffmanCodes = (Map<Byte, String>) ois.readObject();
			// 解码
			byte[] bytes = decode(huffmanCodes, huffmanBytes);
			// 将 bytes 数组写入到目标文件
			os = new FileOutputStream(dstFile);
			os.write(bytes);
		} catch (Exception e) {
			System.out.println(e.getMessage());
		} finally {
			try {
				os.close();
				ois.close();
				is.close();
			} catch (Exception e2) {
				System.out.println(e2.getMessage());
			}
		}
	}

	/**
	 * 数据解压
	 * 思路
	 * 1. 将 huffmanCodeBytes [-88, -65, -56, -65, -56, -65, -55, 77, -57, 6, -24, -14, -117, -4, -60, -90, 28]
	 * 重写先转成 赫夫曼编码对应的二进制的字符串 "1010100010111..."
	 * 2. 赫夫曼编码对应的二进制的字符串 "1010100010111..." →  对照 赫夫曼编码表 →  "i like like like java do you like a java"
	 *
	 * @param huffmanCodes 赫夫曼表
	 * @param huffmanBytes 压缩后的二进制数组
	 * @return 原二进制数组
	 */
	private static byte[] decode(Map<Byte, String> huffmanCodes, byte[] huffmanBytes) {
		// 1. 先得到 huffmanBytes 对应的 二进制的字符串
		StringBuilder stringBuilder = new StringBuilder();
		for (int i = 0; i < huffmanBytes.length; i++) {
			byte b = huffmanBytes[i];
			// 是否是最后一个
			boolean flag = (i == huffmanBytes.length - 1);
			stringBuilder.append(byteToBitString(!flag, b));
		}
		// 把字符串安装指定的赫夫曼编码进行解码
		// 把赫夫曼编码表进行调换，因为反向查询 a->100 100->a
		Map<String, Byte> map = new HashMap<>();
		for (Map.Entry<Byte, String> entry : huffmanCodes.entrySet()) {
			map.put(entry.getValue(), entry.getKey());
		}

		// 创建要给集合，存放 byte
		List<Byte> list = new ArrayList<>();
		for (int i = 0; i < stringBuilder.length(); ) {
			int count = 1;
			boolean flag = true;
			Byte b = null;

			while (flag) {
				//1010100010111...
				String key = stringBuilder.substring(i, i + count);
				b = map.get(key);
				if (b == null) {
					count++;
				} else {
					flag = false;
				}
			}
			list.add(b);
			i += count;
		}
		// 当 for 循环结束后，我们 list 中就存放了所有的字符
		// 把 list 中的数据放入到 byte[] 并返回
		byte b[] = new byte[list.size()];
		for (int i = 0; i < b.length; i++) {
			b[i] = list.get(i);
		}
		return b;
	}

	/**
	 * 将一个 byte 转成一个二进制的字符串
	 *
	 * @param flag 标志是否需要补高位如果是 true（不是最后一个） ，表示需要补高位，如果是 false（最后一个）无需补高位
	 * @param b    传入的 byte
	 * @return 该 b 对应的二进制的字符串(注意是按补码返回)
	 */
	private static String byteToBitString(boolean flag, byte b) {
		int temp = b;
		if (flag) {
			// 如果是正数我们还存在补高位 按位与 256， 1 0000 0000 | 0000 0001 => 1 0000 0001
			temp |= 256;
		}
		String str = Integer.toBinaryString(temp);
		if (flag) {
			return str.substring(str.length() - 8);
		} else {
			return str;
		}
	}

	/**
	 * 赫夫曼压缩
	 *
	 * @param bytes 要压缩的字节数组
	 * @return
	 */
	private static byte[] huffmanZip(byte[] bytes) {
		List<Node> nodes = getNodes(bytes);
		Node huffmanTreeRoot = createHuffmanTree(nodes);
		Map<Byte, String> huffmanCodes = getHuffmanCodes(huffmanTreeRoot);
		return zip(bytes, huffmanCodes);
	}


	/**
	 * 把字节数组根据额对应的赫夫曼表进行压缩
	 *
	 * @param bytes        需要压缩的字节数组
	 * @param huffmanCodes 赫夫曼表
	 * @return
	 */
	private static byte[] zip(byte[] bytes, Map<Byte, String> huffmanCodes) {
		StringBuilder sb = new StringBuilder();
		for (final byte b : bytes) {
			sb.append(huffmanCodes.get(b));
		}
		// 赫夫曼数组的长度，下面等价于：int len = (stringBuilder.length() + 7) / 8;
		int len;
		if (stringBuilder.length() % 8 == 0) {
			len = stringBuilder.length() / 8;
		} else {
			len = stringBuilder.length() / 8 + 1;
		}
		// 创建赫夫曼数组
		byte[] huffmanCodeBytes = new byte[len];
		int index = 0;
		for (int i = 0; i < stringBuilder.length(); i += 8) {
			String strByte;
			if (i + 8 > stringBuilder.length()) {
				strByte = stringBuilder.substring(i);
			} else {
				strByte = stringBuilder.substring(i, i + 8);
			}
			huffmanCodeBytes[index] = (byte) Integer.parseInt(strByte, 2);
			index++;
		}
		return huffmanCodeBytes;
	}

	/**
	 * 生成赫夫曼树
	 */
	public static Node createHuffmanTree(List<Node> nodeList) {
		while (nodeList.size() > 1) {
			Collections.sort(nodeList);
			Node leftNode = nodeList.get(0);
			Node rightNode = nodeList.get(1);
			Node parentNode = new Node(null, leftNode.weight + rightNode.weight);
			parentNode.left = leftNode;
			parentNode.right = rightNode;
			nodeList.remove(leftNode);
			nodeList.remove(rightNode);
			nodeList.add(parentNode);
		}
		return nodeList.get(0);
	}

	/**
	 * 根据赫夫曼树生成对应的赫夫曼表
	 * 类似于：a=100 d=11000 u=11001 e=1110 v=11011 i=101 y=11010 j=0010 k=1111 l=000 o=0011
	 *
	 * @param root 赫夫曼树的根节点
	 * @return 赫夫曼表
	 */
	public static Map<Byte, String> getHuffmanCodes(Node root) {
		if (root == null) {
			return null;
		}
		// 处理左子树
		getCodes(root.left, "0", stringBuilder);
		// 处理右子树
		getCodes(root.right, "1", stringBuilder);
		return huffmanCodes;
	}

	/**
	 * 将传入的 node 结点的所有叶子结点的赫夫曼编码得到，并放入到 huffmanCodes 集合
	 *
	 * @param node 传入节点
	 * @param code 路径: 左子结点是 0, 右子结点 1
	 * @param sb   用于拼接路径
	 */
	private static void getCodes(Node node, final String code, final StringBuilder sb) {
		if (node == null) {
			return;
		}
		StringBuilder sb2 = new StringBuilder(sb);
		// 拼接路径
		sb2.append(code);
		// 非叶子节点，继续递归处理
		if (node.data == null) {
			getCodes(node.left, "0", sb2);
			getCodes(node.right, "1", sb2);
		} else {
			// 叶子节点
			huffmanCodes.put(node.data, sb2.toString());
		}
	}

	/**
	 * 字节数组转为 node<byte,出现次数> 节点放入 list
	 *
	 * @param bytes 字节数组
	 * @return List [Node[date=97 ,weight = 5], Node[date=32,weight = 9]......],
	 */
	private static List<Node> getNodes(byte[] bytes) {
		List<Node> nodes = new ArrayList<>();
		Map<Byte, Integer> counts = new HashMap<>();
		for (byte b : bytes) {
			// 出现次数
			Integer count = counts.get(b);
			if (count == null) {
				counts.put(b, 1);
			} else {
				counts.put(b, count + 1);
			}
		}
		for (Map.Entry<Byte, Integer> entry : counts.entrySet()) {
			nodes.add(new Node(entry.getKey(), entry.getValue()));
		}
		return nodes;
	}

	/**
	 * 前序遍历
	 */
	private static void preOrder(Node root) {
		if (root != null) {
			root.preOrder();
		}
	}
}

class Node implements Comparable<Node> {
	Byte data; // 'a' => 97 ' ' => 32
	int weight;// 出现次数
	Node left;
	Node right;

	public Node(Byte data, int weight) {

		this.data = data;
		this.weight = weight;
	}

	@Override
	public int compareTo(Node o) {
		return this.weight - o.weight;
	}

	public String toString() {
		return "Node [data = " + data + " weight=" + weight + "]";
	}

	/**
	 * 前序遍历
	 */
	public void preOrder() {
		System.out.println(this);
		if (this.left != null) {
			this.left.preOrder();
		}
		if (this.right != null) {
			this.right.preOrder();
		}
	}
}
